from data_loader_2D import load_dataset
import numpy as np
import time
from rknn_dpc import RKNN_DPC
from params import DATASET_PARAMS
from draw_pictures import visualize_results

if __name__ == "__main__":

    # 多个数据集
    datasets_to_compare = [
        'Jain',
        'Flame',
        'Spiral',
        # 'Compound',
        'Aggregation',
        'D31'
    ]

    for dataset_name in datasets_to_compare:
        X, true_labels = load_dataset(dataset_name)
        if X is not None:
            # 获取数据集对应的参数
            params = DATASET_PARAMS.get(dataset_name)
            print(f"参数设置: k={params['k']}, m={params['m']}, R={params['R']}")

            start_time = time.time()

            # 应用 RKNN-DPC 算法
            rknn_dpc = RKNN_DPC(k=params['k'], m=params['m'], R=params['R'])
            rknn_dpc.rknn_dpc(X)
            labels = rknn_dpc.predict(X)

            print(f"聚类完成，用时: {time.time() - start_time:.2f}秒")
            print(f"类簇中心索引: {rknn_dpc.centers}")
            print(f"类簇数量: {len(np.unique(labels))}")

            # 可视化结果
            visualize_results(dataset_name, X, labels, true_labels, rknn_dpc, params['k'], params['R'])

            # print("labels:", labels)

            print("--------------------------------------------------------------------------------------------")